import json
from openai import OpenAI
import base64
import fitz  # PyMuPDF
import os
import re
import time

def get_image_content(image_path):
    """将图片转换为base64编码字符串"""
    try:
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
    except FileNotFoundError:
        print(f"错误：文件 {image_path} 未找到")
        return None

def call_lmstudio_api(image_base64, article_title):
    """调用LM Studio API（使用OpenAI兼容模式）"""
    client = OpenAI(
        base_url="http://localhost:1234/v1",
        api_key="lm-studio"
    )

    try:
        completion = client.chat.completions.create(
            model="gemma-3-27b-it",
            messages=[
                {"role": "system", "content": f"""你必须是严格遵循JSON格式的问答对生成器。根据图片中的《{article_title}》文章内容生成5-10个严谨问答对，输出格式：
{{
  "qa_pairs": [
    {{"question": "...", "answer": "..."}}
  ]
}}
硬性要求：
1.问题必须直接来自文章核心内容，聚焦"{article_title}"主题
2.问题必须包含与文章标题相关的术语
3.答案必须精确对应原文，使用陈述句
4.严格保持JSON结构，禁止注释或额外字段
示例：
{{
  "qa_pairs": [
    {{"question": "{article_title.split('：')[0]}的核心价值是什么？", "answer": "..."}}
  ]
}}"""},
                {"role": "user", "content": [
                    {"type": "text", "text": "请严格按以下JSON格式根据图片内容生成5-10个严谨的问答对：\n{\n  \"qa_pairs\": [\n    {\"question\": \"兵棋推演的核心目的是什么？\", \"answer\": \"通过模拟检验作战方案、发现问题并改进\"}\n  ]\n}\n要求：\n1.问题直接明确，不要包含'根据文章'、'文章中提到'等无关内容\n2.答案使用简洁陈述句，避免引用格式说明"},
                    {"type": "image_url", "image_url": {
                        "url": f"data:image/jpeg;base64,{image_base64}"  # 修正base64分隔符
                    }}
                ]}
            ],
            temperature=0.0,
            max_tokens=1024
        )
        time.sleep(10)  # 添加10秒延时缓解显卡压力
        return completion.choices[0].message.content
    except Exception as e:
        print(f"API请求失败: {e}")
        return None

def process_images_to_json(image_dir="temp-jpg", output_dir="out-temp"):
    """处理图片生成JSON问答对"""
    try:
        os.makedirs(output_dir, exist_ok=True)
        qa_data = []
        
        for img_file in os.listdir(image_dir):
            if img_file.lower().endswith(('.jpeg', '.jpg')):
                img_path = os.path.join(image_dir, img_file)
                print(f"\n处理图片: {img_file}")
                
                if img_base64 := get_image_content(img_path):
                    # 从图片文件名提取文章标题（去除页码和扩展名）
                    article_title = re.sub(r'_?\d{3}\.jpe?g$', '', img_file, flags=re.IGNORECASE)
                    if response := call_lmstudio_api(img_base64, article_title):
                        try:
                            # 简化版JSON清洗
                            cleaned_response = re.sub(r'[“”]', '"', response)  # 统一中文引号为英文双引号
                            # 提取最外层JSON结构
                            json_match = re.search(r'\{[\s\S]*\}', cleaned_response)
                            if json_match:
                                cleaned_response = json_match.group(0)
                            # 基础语法修复
                            cleaned_response = re.sub(r'(?<=\})\s*(?=\{)', ',', cleaned_response)  # 对象间添加逗号
                            cleaned_response = re.sub(r'(?<=[}\]"])\s*(?=["{[])', ',', cleaned_response)  # 数组元素间添加逗号
                            
                            try:
                                # 添加严格的JSON格式验证
                                json_data = json.loads(cleaned_response)
                                # 从文章标题生成动态关键词
                                # 优化标题处理：保留完整术语组合
                                title_keywords = re.split(r'[：_\-()（）]', article_title)[0]
                                # 提取连续中文词汇组合（2-6字）
                                keywords = re.findall(r'[\u4e00-\u9fa5]{2,6}', title_keywords)
                                # 添加完整标题作为后备关键词
                                keywords.append(title_keywords.strip())
                                
                                if not keywords:
                                    print(f"警告：无法从标题'{article_title}'提取有效关键词，跳过内容验证")
                                    keywords = [article_title]  # 使用完整标题作为后备
                                
                                valid_pairs = []
                                for pair in json_data.get('qa_pairs', []):
                                    # 检查字段格式
                                    if not (isinstance(pair.get('question'), str) and 
                                           isinstance(pair.get('answer'), str)):
                                        continue
                                    
                                    # 检查问题包含至少一个标题关键词
                                 #   question_text = pair['question']
                                 #   if not any(keyword in question_text for keyword in keywords):
                                 #       print(f"过滤无关问题（缺少关键词）: {question_text}")
                                 #       continue
                                    
                                    valid_pairs.append(pair)
                                
                                if not valid_pairs:
                                    print(f"警告：图片{img_file}没有有效的问答对（需包含关键词：{'/'.join(keywords)}）")
                                
                                json_data["qa_pairs"] = valid_pairs
                            except json.JSONDecodeError as e:
                                print(f"JSON解析失败，清洗后的内容：\n{cleaned_response}")
                                continue  # 跳过无效的JSON响应
                            
                            if "qa_pairs" in json_data and isinstance(json_data["qa_pairs"], list):
                                # 添加来源信息
                                json_data["source_image"] = img_file
                                # 生成输出文件名
                                output_file = os.path.splitext(img_file)[0] + ".json"
                                output_path = os.path.join(output_dir, output_file)
                                # 保存JSON文件
                                with open(output_path, 'w', encoding='utf-8') as f:
                                    json.dump(json_data, f, ensure_ascii=False, indent=2)
                                print(f"已保存: {output_path}")
                                qa_data.extend(json_data["qa_pairs"])
                            else:
                                print("响应格式无效，缺少qa_pairs字段")
                        except json.JSONDecodeError:
                            print("API响应不是有效的JSON格式")
        return qa_data
    except Exception as e:
        print(f"处理图片失败: {e}")
        return []

def merge_json_files(input_dir="out-temp", output_dir="output"):
    """合并多个JSON文件"""
    try:
        os.makedirs(output_dir, exist_ok=True)
        results = {
            "total_entries": 0,
            "sources": [],
            "qa_pairs": []
        }
        
        for json_file in os.listdir(input_dir):
            if json_file.endswith('.json'):
                file_path = os.path.join(input_dir, json_file)
                with open(file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    # 收集元数据
                    results["sources"].append({
                        "file_name": json_file,
                        "source_image": data.get("source_image", ""),
                        "qa_count": len(data.get("qa_pairs", []))
                    })
                    # 合并问答对
                    results["qa_pairs"].extend(data.get("qa_pairs", []))
        
        # 更新统计信息
        results["total_entries"] = len(results["qa_pairs"])
        
        output_path = os.path.join(output_dir, "combined_output.json")
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        
        print(f"\n合并完成！最终输出文件: {output_path}")
        return True
    except Exception as e:
        print(f"合并文件失败: {e}")
        return False

def convert_pdf_to_jpeg(pdf_path, output_dir="temp-jpg"):
    """将PDF转换为JPEG图片"""
    try:
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 从文件名获取标题（保留中文，仅去除特殊符号）
        base_name = os.path.basename(pdf_path)
        title = re.sub(r'[\\/:*?"<>|]', '', os.path.splitext(base_name)[0])  # 只去除文件系统非法字符
        
        # 打开PDF文件
        doc = fitz.open(pdf_path)
        
        for page_num in range(len(doc)):
            page = doc.load_page(page_num)
            pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))  # 提高分辨率到300dpi
            output_path = os.path.join(
                output_dir, 
                f"{title}_{page_num+1:03d}.jpeg"
            )
            pix.save(output_path)
            print(f"已生成：{output_path}")
            
        return True
    except Exception as e:
        print(f"PDF转换失败：{e}")
        return False

if __name__ == "__main__":
    # 处理input目录下的所有PDF文件
    input_dir = "input"
    for filename in os.listdir(input_dir):
        if filename.lower().endswith(".pdf"):
            pdf_path = os.path.join(input_dir, filename)
            print(f"\n正在处理：{filename}")
            if convert_pdf_to_jpeg(pdf_path):
                print("转换成功！")
                
                # 处理生成的图片
                print("\n开始处理图片生成问答对...")
                if process_images_to_json():
                    # 合并JSON文件
                    print("\n开始合并JSON文件...")
                    if merge_json_files():
                        print("全部流程完成！")
                    else:
                        print("JSON合并失败")
                else:
                    print("问答对生成失败")
            else:
                print("转换失败")
